use axum::extract::ConnectInfo;
use axum::http::{Request, Response};
use std::fmt::{Display, Formatter};
use std::net::SocketAddr;
use std::sync::LazyLock;
use std::time::Duration;
use tower_http::trace::{DefaultOnBodyChunk, DefaultOnEos, MakeSpan, OnRequest, OnResponse};
use tower_http::trace::{HttpMakeClassifier, TraceLayer};
use tracing::Span;
use crate::util::xid;

type CustomTraceLayer = TraceLayer<
    HttpMakeClassifier,
    MyMakeSpan,
    StartOnRequest,
    LatencyOnResponse,
    DefaultOnBodyChunk,
    DefaultOnEos,
    (),
>;
static TRACE_LAYER: LazyLock<CustomTraceLayer> = LazyLock::new(|| {
    TraceLayer::new_for_http()
        .make_span_with(MyMakeSpan)
        .on_request(StartOnRequest)
        .on_failure(())
        .on_response(LatencyOnResponse)
});

#[derive(Clone)]
pub struct StartOnRequest;

impl<B> OnRequest<B> for StartOnRequest {
    fn on_request(&mut self, _: &Request<B>, _: &Span) {
        tracing::info!("请求开始");
    }
}

#[derive(Clone)]
pub struct LatencyOnResponse;

impl<B> OnResponse<B> for LatencyOnResponse {
    fn on_response(self, _: &Response<B>, latency: Duration, _: &Span) {
        // tracing::info!(
        //     耗时 = %Latency(latency),
        //     // status = response.status().as_u16(),
        //     "请求结束"
        // );

        tracing::info!(
            "请求结束, 耗时: {}", Latency(latency)
        );
    }
}

struct Latency(Duration);

impl Display for Latency {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        if self.0.as_millis() > 0 {
            write!(f, "{}ms", self.0.as_millis())
        } else {
            write!(f, "{}μs", self.0.as_micros())
        }
    }
}

#[derive(Clone)]
pub struct MyMakeSpan;

impl<B> MakeSpan<B> for MyMakeSpan {
    fn make_span(&mut self, request: &Request<B>) -> Span {
        let uri = request.uri();
        let id = xid();
        let extensions = request.extensions();
        let connect = extensions.get::<ConnectInfo<SocketAddr>>();
        if let Some(ConnectInfo(addr)) = connect {
            tracing::info_span!("Request", id = %id, ip = %addr.ip(), uri = %uri)
        } else {
            tracing::info_span!("Request", id = %id, uri = %uri)
        }
    }
}

/// 日志中间件
pub fn get_trace_layer() -> &'static CustomTraceLayer {
    &TRACE_LAYER
}
